#include "QFeatureTrackingDemoApplication.hpp"
#include <boost/function_output_iterator.hpp>

const double QFeatureTrackingDemoApplication::MIN_MATCH_SCORE = 0.65;

QFeatureTrackingDemoApplication::QFeatureTrackingDemoApplication(int& argc, char** argv) :
   QApplication(argc, argv), _frameIndex(0), _detector(10, true)
{
   _pVideoSource = new btl::extra::VideoSourceUEye("/home/pgrad/glynn/shared/btl2/uEye.ini");
   _pTimer = new QTimer;
   _pView = new btl::extra::QImageViewRGB;
   _pView->show();
   connect(_pTimer, SIGNAL(timeout()), this, SLOT(update()));
   _pTimer->start(30);
}

QFeatureTrackingDemoApplication::~QFeatureTrackingDemoApplication() {}

void QFeatureTrackingDemoApplication::update()
{
   btl::ImageRegionConstRGB frame = _pVideoSource->getNextFrame();
   btl::ImageGrey grey;
   grey.allocate(frame.width(), frame.height()); // FIXME: need to improve interface for image colour conversion
   btl::convert_image(grey, frame);

   _displayImage.clone(frame);
   _painter.setImage(_displayImage);

   trackFeatures(grey);
   dropDeadFeatures();
   initialiseNewFeatures(grey);
   visualiseFeatures();

   _pView->setImage(_displayImage);

   ++_frameIndex;
}

void QFeatureTrackingDemoApplication::trackFeatures(btl::ImageRegionConstGrey frame)
{
   _detectedFeatures.clear();

   // detect FAST features to try to find features we're tracking
   std::vector<cv::KeyPoint> cvKeyPoints;
   _detector.detect(btl::cvUnsafeUnconstImage(frame), cvKeyPoints);
   const int halfPatchSize = PATCH_SIZE / 2;
   Eigen::AlignedBox<int, 2> safeRegion(
      Eigen::Vector2i(halfPatchSize, halfPatchSize),
      Eigen::Vector2i(frame.width() - halfPatchSize - 1, frame.height() - halfPatchSize - 1));
   std::vector<cv::KeyPoint>::iterator it;
   for(it = cvKeyPoints.begin(); it != cvKeyPoints.end(); ++it)
   {
      cv::KeyPoint cvKeyPoint = *it;
      int x = cvKeyPoint.pt.x;
      int y = cvKeyPoint.pt.y;
      if(safeRegion.contains(Eigen::Vector2i(x, y)))
         _detectedFeatures.push_back(
            DetectedFeature(
               Eigen::Vector2f(float(x), float(y)),
               btl::PatchFeatureDescriptor<btl::PixelGrey, PATCH_SIZE, PATCH_SIZE>(frame, x, y)));

   }

   btl::married_matches(
      _trackedFeatures.begin(), _trackedFeatures.end(),
      boost::make_function_output_iterator(FeatureTrackUpdater(_frameIndex)),
      btl::fixed_radius_searcher(_detectedFeatures.begin(), _detectedFeatures.end(), MAX_MATCH_RADIUS),
      ScoreNCC(), MIN_MATCH_SCORE);
}

void QFeatureTrackingDemoApplication::dropDeadFeatures()
{
   std::vector<TrackedFeature>::iterator newEnd
   = std::remove_if(_trackedFeatures.begin(), _trackedFeatures.end(),
                    TrackedFeature::IsOlderThan(_frameIndex - MAX_FEATURE_AGE));
   size_t numRemoved = std::distance(newEnd, _trackedFeatures.end());
   size_t numKept = std::distance(_trackedFeatures.begin(), newEnd);
   assert(numRemoved + numKept == _trackedFeatures.size());
   std::cout << "Removing " << numRemoved << " features (keeping " << numKept << ").\n";
   _trackedFeatures.erase(newEnd, _trackedFeatures.end());
}

void QFeatureTrackingDemoApplication::initialiseNewFeatures(btl::ImageRegionConstGrey frame)
{
   const int FEATURE_MASK_RADIUS = 4;
   const int NUM_GRID_CELLS_X = 10;
   const int NUM_GRID_CELLS_Y = 8;
   int gridCellWidth = frame.width() / NUM_GRID_CELLS_X;
   int gridCellHeight = frame.height() / NUM_GRID_CELLS_Y;
   const int halfPatchSize = PATCH_SIZE / 2;
   Eigen::AlignedBox<int, 2> safeRegion(
      Eigen::Vector2i(halfPatchSize, halfPatchSize),
      Eigen::Vector2i(frame.width() - halfPatchSize - 1, frame.height() - halfPatchSize - 1));
   int numFeaturesPerCell[NUM_GRID_CELLS_Y][NUM_GRID_CELLS_X];
   for(int i = 0; i < NUM_GRID_CELLS_Y; ++i)
      for(int j = 0; j < NUM_GRID_CELLS_X; ++j)
         numFeaturesPerCell[i][j] = 0;

   btl::Image<btl::PixelGrey> mask(frame.width(), frame.height(), 0xff);
   btl::extra::ImagePainter<btl::PixelGrey> maskpainter(mask);
   std::vector<TrackedFeature>::iterator it;
   for(it = _trackedFeatures.begin(); it != _trackedFeatures.end(); ++it)
   {
      TrackedFeature feature = *it;
      Eigen::Vector2f pos = feature.position();
      maskpainter.drawFilledCircle(pos.cast<int>(), FEATURE_MASK_RADIUS, 0);
      //mask[(int)pos(1)][(int)pos(0)] = 0;
      int gridCellX = (int)pos(0) / gridCellWidth;
      int gridCellY = (int)pos(1) / gridCellHeight;
      ++numFeaturesPerCell[gridCellY][gridCellX];
   }

   for(int i = 0; i < NUM_GRID_CELLS_Y; ++i)
   {
      for(int j = 0; j < NUM_GRID_CELLS_X; ++j)
      {
         int topLeftY = gridCellHeight * i;
         int topLeftX = gridCellWidth * j;
         btl::ImageRegionConstGrey region = frame.region(topLeftX, topLeftY, gridCellWidth, gridCellHeight);

         std::vector<cv::KeyPoint> cvKeyPoints;
         btl::ImageRegionGrey maskregion = mask.region(topLeftX, topLeftY, gridCellWidth, gridCellHeight);
         cv::Mat cvmask = btl::cvImage(maskregion);
         _detector.detect(btl::cvUnsafeUnconstImage(region), cvKeyPoints, cvmask);
         std::vector<cv::KeyPoint>::iterator it;
         for(it = cvKeyPoints.begin(); it != cvKeyPoints.end(); ++it)
         {
            cv::KeyPoint cvKeyPoint = *it;
            int x = cvKeyPoint.pt.x + topLeftX;
            int y = cvKeyPoint.pt.y + topLeftY;
            if(safeRegion.contains(Eigen::Vector2i(x, y)))
               _trackedFeatures.push_back(
                  TrackedFeature(
                     _frameIndex,
                     Eigen::Vector2f(x, y),
                     btl::PatchFeatureDescriptor<btl::PixelGrey, PATCH_SIZE, PATCH_SIZE>(frame, x, y)));
         }
      }
   }
}

void QFeatureTrackingDemoApplication::visualiseFeatures()
{
   std::vector<TrackedFeature>::const_iterator it;
   btl::SmallImage<btl::PixelRGB, PATCH_SIZE, PATCH_SIZE> colourpatch;
   for(it = _trackedFeatures.begin(); it != _trackedFeatures.end(); ++it)
   {
      Eigen::Vector2i pos = it->position().cast<int>();
      bool newFeature = (it->firstSeenAt() == _frameIndex);
      bool trackedFeature = (it->lastSeenAt() == _frameIndex);
      //display colour coded patch
      for(int i = 0; i < PATCH_SIZE; ++i)
      {
         for(int j = 0; j < PATCH_SIZE; ++j)
         {
            btl::PixelGrey patchPixel = it->descriptor().patch()[i][j];
            btl::PixelRGB newPixel(0, 0, 0);
            if(newFeature)
               newPixel.blue() = patchPixel;
            else if(trackedFeature)
               newPixel.green() = patchPixel;
            else
               newPixel.red() = patchPixel;
            colourpatch[i][j] = newPixel;
         }
      }
      Eigen::Vector2i topLeft = pos - Eigen::Vector2i(PATCH_SIZE / 2, PATCH_SIZE / 2);
      _displayImage.region(topLeft(0), topLeft(1), PATCH_SIZE, PATCH_SIZE).copy(colourpatch);
      //draw line trace between current patch position and initial position
      if(!newFeature)
      {
         btl::PixelRGB colour;
         if(trackedFeature) colour = btl::PixelRGB(0, 255, 0);
         else colour = btl::PixelRGB(255, 0, 0);
         Eigen::Vector2i firstPos = it->firstPosition().cast<int>();
         _painter.drawLine(firstPos, pos, colour);
      }

   }

   //for (std::vector<keypoint_type>::const_iterator
   //      it = _detectedFeatures.begin(); it != _detectedFeatures
}
